# plot_rivers.py  

"""This is the driver code for getting river data.

2/4/2014  Parker MacCready
"""

import numpy as np

from pandas import Series, DataFrame
import pandas as pd

import river_class
reload(river_class) # needed while editing class definition

import Rfun2
reload(Rfun2)

Ldir = Rfun2.get_ldir()

# make a dict of station names (key = USGS number)
inlist = pd.read_csv('./Files_USGS/USGS_good_plus.csv') # a pandas DataFrame
sta_num = inlist['Station Number'].values # a list of ints
sta_name = inlist['Name'].values # a list of strings

if False:
    # use a list all the rivers
    riv_list_to_get = sta_name
else:
    # or just some
    riv_list_to_get = ['skagit',  'snohomish', 'deschutes']

id_dict = {}    
for key, value in zip(sta_name, sta_num):
    if key in riv_list_to_get:
        id_dict[key] = value
    
# loop over all the rivers in the dict, and store the result in a dict
riv_dict = {}
riv_dict_raw = {}
good_riv_list = []
bad_riv_list = []
for riv_name in id_dict.keys():
    
    # set the id number
    id = id_dict[riv_name]
         
    # create a river object
    riv = river_class.River(id, riv_name)

    # add the data to the object
    riv.get_data()

    # next make daily averages and put all rivers in a single array

    # pack the data as a pandas Series
    if len(riv.T) > 0:
        qt = Series(riv.Q, index=riv.T)
        # change the time to UTC (a new Series)
        qt_utc = qt.tz_convert('UTC')
        # form daily means (a new Series)
        qt_rs = qt_utc.resample('D', how='mean', label='right', loffset='-12h')    
        # make a dict of Series
        riv_dict[riv_name] = qt_rs
        riv_dict_raw[riv_name] = qt_utc
        good_riv_list.append(riv_name)
    else:
        bad_riv_list.append(riv_name)

# convert riv_dict to a DataFrame
riv_df = DataFrame(riv_dict)
riv_df_raw = DataFrame(riv_dict_raw)
# make sure all values are filled (should use climatology)
riv_df = riv_df.fillna(method='bfill')
riv_df = riv_df.fillna(method='ffill') 
    
# add a column of seconds since 1970/01/01
tt = riv_df.index   
TT = []
for ttt in tt:    
    TT.append(ttt.value)
T = np.array(TT)
Tsec = T/1e9
riv_df['Tsec1970'] = Tsec
# and a column of yearday
riv_df['yearday'] = tt.dayofyear

# Writing and Reading
#
# first name the index
iname = 'DateTimeUTC'
riv_df.index.name = iname
# now write it out to csv
riv_df.to_csv('./output/riv_df.csv')
# and this reads back in the exact same DataFrame, proving it is good for storage
riv_df2 = pd.read_csv('./output/riv_df.csv', index_col = iname)
#
# this is designed to be easier for Matlab to open
riv_df.to_csv('./output/riv_df_for_matlab.csv', index=False)

# plotting
do_plot = True        
if do_plot:
    riv_list_to_plot = good_riv_list
    Rfun.plot_rivers(riv_df, riv_df_raw, riv_list_to_plot)


    
    
    
    
    

